from backtrader import indicator
import pandas as pd
from pymongo import MongoClient
import datetime
import time
import math
import backtrader as bt
import backtrader.indicators as btind
import backtrader.feeds as btfeeds
import talib
from my_plot import *
from py_mysql import *

start =time.time()
class PandasData_more(bt.feeds.PandasData):
    lines = ('m_RSI','h_RSI') # 要添加的线
    # 设置 line 在数据源上的列位置
    params=(
        ('m_RSI', 5),
        ('h_RSI', 7),
           ) 

# 创建策略继承bt.Strategy
class TestStrategy(bt.Strategy):
    params = (
        # 均线参数设置15天，15日均线
        ('maperiod', 15),
    )

    def log(self, txt, dt=None):
        # 记录策略的执行日志
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # 保存收盘价的引用
        self.dataclose = self.datas[0].close
        # 跟踪挂单
        self.order = None
        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None

        self.RB_h_Rsi = self.datas[0].h_RSI  #主期货大周期对RSI
        self.RB_m_Rsi = self.datas[0].m_RSI  #主期货小周期对RSI

        self.I_h_Rsi = self.datas[1].h_RSI  #对冲期货大周期对RSI
        self.I_m_Rsi = self.datas[1].m_RSI  #对冲期货小周期对RSI

        self.maxRsi = 60  #最大偏离度
        self.minRsi = 10    #最小偏离度

        # 23:00，15:00  该时段平仓
        self.close_time_arr = ['23:00:00']

    def next(self):
        self.h_Rsi = self.RB_h_Rsi[0] - self.I_h_Rsi[0]  #大周期偏离度
        self.m_Rsi = self.RB_m_Rsi[0] - self.I_m_Rsi[0]  #小周期偏离度
        
        self.RB_size = self.getposition(self.datas[0]).size  #主期货交易量
        self.I_size = self.getposition(self.datas[1]).size   #对冲期货交易量
        # print(self.RB_size)

        self.time = str(bt.num2date(self.lines.datetime[0]))

        if self.time.split(' ')[1] in self.close_time_arr:
            self.close()
        else:
            if (self.h_Rsi > 0 and self.h_Rsi > self.minRsi and self.h_Rsi < self.maxRsi and self.m_Rsi > self.minRsi ):
                if(self.RB_size > -40):
                    self.sell(data = self.datas[0],size = 40)

                if(self.I_size < 300):
                    self.buy(data = self.datas[1],size = 300)

                
            if (self.h_Rsi < 0 and 
            abs(self.h_Rsi) > self.minRsi and 
            abs(self.h_Rsi) < self.maxRsi and
            self.m_Rsi < 0 and
            abs(self.m_Rsi) > self.minRsi):

                if(self.RB_size < 40):
                    self.buy(data = self.datas[0],size = 40)
                    
                if(self.I_size > -300):
                    self.sell(data = self.datas[1],size = 300)





     # 订单状态通知，买入卖出都是下单
    def notify_order(self, order):
        # 提交了/接受了,  买/卖订单什么都不做
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 检查一个订单是否完成
        # 注意:当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm

                print({'状态':'买入','date': self.time, 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
            elif order.issell():
                print({'状态':'卖出','date': self.time, 'value': round(order.executed.price, 4), 'size': order.executed.size, 'comm': round(order.executed.comm, 4)})
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            print('订单取消/保证金不足/拒绝')
        # 其他状态记录为：无法挂单
        self.order = None


    # 交易状态通知，一买一卖算交易（交易净利润）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return


def main(codeArr, startDate, endDate, startcash=10000000, com=0.00002, qts=13):

    # 创建主控制器
    cerebro = bt.Cerebro()
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)

    # 获取数据
    query_db = Mysql_search()
    df = query_db.get_two(codeArr,startDate,endDate)
    for item in df:
        df2 = df[item]
        df2.index = pd.to_datetime( df2.index, utc= True)

        df2['m_RSI'] = talib.RSI(df2.close, timeperiod=28)
        df2['ma60'] = talib.MA(df2.close,60)
        df2['h_RSI'] = talib.RSI(df2.ma60, timeperiod=28*60)
        data = PandasData_more(dataname = df2,timeframe=bt.TimeFrame.Minutes)
        cerebro.adddata(data)

    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    cerebro.broker.set_coc(True)
    
    print('期初总资金: %.2f' % cerebro.broker.getvalue())

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name= '_TimeReturn')
    result = cerebro.run()

    print('期末总资金: %.2f' % cerebro.broker.getvalue())

    # custom_plot(result)
    cerebro.plot()


if __name__ == '__main__':
    startDate = '2019-01-01'  #开始时间
    endDate = '2019-12-31'   #结束时间
    codeArr = ['RB','I']   #品种数组
    
    main(codeArr,startDate,endDate)

end = time.time()
print('运行总时长: %s Seconds'%(end-start))